"""
-*- coding: utf-8 -*-
@File  : generate_flow_sh.py
@author: ZhenyuYang
@Time  : 2023/03/31 11:14
"""
import os

with open("generate_flow.sh", "w") as f:
    f.write("#!/usr/bin/env bash\n")
    f.write("cd ../tools/gmflow-main\n")
raw_frame_path = "../data/flow_prepare_rawframes"
dataset = ["casme2", "smic", "samm"]
labels = ["negative", "positive", "surprise"]
if not os.path.exists("../data/flows_raw"):
    os.mkdir("../data/flows_raw")
for label in labels:
    if not os.path.exists("../data/flows_raw/{}".format(label)):
        os.mkdir("../data/flows_raw/{}".format(label))
    raw_frame_class_path = "{}/{}".format(raw_frame_path, label)
    folder_list = os.listdir(raw_frame_class_path)
    for folder in folder_list:
        folder_dataset = folder.split("-")[0]
        if folder_dataset not in dataset:
            continue
        src_path = os.path.abspath("{}/{}".format(raw_frame_class_path, folder))
        dst_path = os.path.abspath("../data/flows_raw/{}/{}".format(label, folder))
        # src_path = "{}/{}".format(raw_frame_class_path, folder)
        # dst_path = "../data/flows/{}/{}".format(label, folder)
        with open("generate_flow.sh", "a") as f:
            f.writelines(["CUDA_VISIBLE_DEVICES=0 python main.py ",
                          "--inference_dir {} ".format(src_path),
                          "--output_path {} ".format(dst_path),
                          "--resume pretrained/gmflow_sintel-0c07dcb3.pth"])
            f.write("\t#\r\n")
